package org.xserver.component.listener;

import java.util.Collections;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;

import org.jboss.netty.handler.execution.ChannelDownstreamEventRunnable;
import org.jboss.netty.handler.execution.ChannelUpstreamEventRunnable;
import org.jboss.netty.handler.execution.OrderedMemoryAwareThreadPoolExecutor;
import org.jboss.netty.util.ObjectSizeEstimator;

public class XServerMemoryAwareThreadPoolExecutor extends
		OrderedMemoryAwareThreadPoolExecutor {
	private final Map<Runnable, Boolean> inProgress = new ConcurrentHashMap<Runnable, Boolean>();
	private final ThreadLocal<Long> startTime = new ThreadLocal<Long>();
	private long totalTime;
	private int totalTasks;

	public XServerMemoryAwareThreadPoolExecutor(int corePoolSize,
			long maxChannelMemorySize, long maxTotalMemorySize) {
		super(corePoolSize, maxChannelMemorySize, maxTotalMemorySize);
	}

	public XServerMemoryAwareThreadPoolExecutor(int corePoolSize,
			long maxChannelMemorySize, long maxTotalMemorySize,
			long keepAliveTime, TimeUnit unit) {
		super(corePoolSize, maxChannelMemorySize, maxTotalMemorySize,
				keepAliveTime, unit);
	}

	public XServerMemoryAwareThreadPoolExecutor(int corePoolSize,
			long maxChannelMemorySize, long maxTotalMemorySize,
			long keepAliveTime, TimeUnit unit,
			ObjectSizeEstimator objectSizeEstimator, ThreadFactory threadFactory) {
		super(corePoolSize, maxChannelMemorySize, maxTotalMemorySize,
				keepAliveTime, unit, objectSizeEstimator, threadFactory);
	}

	public XServerMemoryAwareThreadPoolExecutor(int corePoolSize,
			long maxChannelMemorySize, long maxTotalMemorySize,
			long keepAliveTime, TimeUnit unit, ThreadFactory threadFactory) {
		super(corePoolSize, maxChannelMemorySize, maxTotalMemorySize,
				keepAliveTime, unit, threadFactory);
	}

	@Override
	protected void beforeExecute(Thread t, Runnable r) {
		super.beforeExecute(t, r);
		if (r instanceof ChannelUpstreamEventRunnable
				|| r instanceof ChannelDownstreamEventRunnable) {
			inProgress.put(r, Boolean.TRUE);
		}
		startTime.set(new Long(System.currentTimeMillis()));
	}

	@Override
	protected void afterExecute(Runnable r, Throwable t) {
		long time = System.currentTimeMillis() - startTime.get().longValue();
		synchronized (this) {
			totalTime += time;
			++totalTasks;
		}
		inProgress.remove(r);
		super.afterExecute(r, t);
	}

	public Set<Runnable> getInProcessTasks() {
		return Collections.unmodifiableSet(inProgress.keySet());
	}

	public synchronized int getTotalTasks() {
		return totalTasks;
	}

	public synchronized double getAverageTaskTime() {
		return (totalTasks == 0) ? 0 : totalTime / totalTasks;
	}
}
